import torch
import numpy as np

# 读取数据
data = torch.load("dataset.pt", weights_only=False)

inconsistent_samples = []  # 用来记录不一致的样本索引
cnt = 0
for i, sample in enumerate(data):
    
    src = sample["src"] if isinstance(sample, dict) else sample[0]
    tgt = sample["tgt"] if isinstance(sample, dict) else sample[1]

    src = np.array(src)
    tgt = np.array(tgt)

    if src.shape[0] != tgt.shape[0]:
        cnt += 1
        inconsistent_samples.append((i, src.shape, tgt.shape))


# 输出结果
if inconsistent_samples:
    print("❌ 以下样本的 src 与 tgt 长度不一致：")
    for idx, src_shape, tgt_shape in inconsistent_samples:
        print(f"样本 {idx}: src shape={src_shape}, tgt shape={tgt_shape}")
else:
    print("✅ 所有样本的 src 与 tgt 长度一致！")

print("数据集大小:", cnt)